from typing import Any, Sequence

import chainlit as cl
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import BaseMessage
from langchain_core.prompt_values import PromptValue
from langchain_core.runnables import Runnable


def find_mcp_name_for_tool(tool_name) -> str | None:
    """
    通过工具名称查询MCP的名称， 默认返回第1个符合要求的
    TODO 存在工具名称相同，但是在不同的MCP中
    :param tool_name: 工具
    :return: str
    """
    if tool_name is None:
        return None
    mcp_tools = cl.user_session.get("mcp_tools", {})
    all_tools = [tool for connection_tools in mcp_tools.values() for tool in connection_tools]
    result = [item for item in all_tools if item["name"] == tool_name]
    if len(result) == 0:
        return None
    return result[0].get("mcp-name")


def add_mcp_tool_to_model(chat_model: BaseChatModel) -> Runnable[
    PromptValue | str | Sequence[BaseMessage | list[str] | tuple[str, str] | str | dict[str, Any]], BaseMessage]:
    """
    添加mcp 工具到指定的chat model中
    :param chat_model: LLM
    :return: chat_model
    """
    mcp_tools = cl.user_session.get("mcp_tools", {})
    all_tools = [tool for connection_tools in mcp_tools.values() for tool in connection_tools]
    return chat_model.bind_tools(tools=all_tools)
